import sqlite3
import os
import pandas as pd
import utils.configFinRAG as configFinRAG
from FinSQL_01_generate import generate_sql
from FinSQL_02_query import query_db
from FinSQL_03_answer_from_SQL import generate_answer
from utils.instances import TOKENIZER, LLM

# 初始化全局变量
g_example_question_list = []
g_example_sql_list = []
g_example_fa_list = []
g_example_info_list = []
g_example_token_list = []

# 封装读取示例数据的函数
def load_example_data():
    global g_example_question_list, g_example_sql_list, g_example_fa_list, g_example_info_list, g_example_token_list
    try:
        sql_examples_file = pd.read_csv(configFinRAG.sql_examples_path, delimiter=",", header=0)
        for _, row in sql_examples_file.iterrows():
            g_example_question_list.append(row['问题'])
            g_example_sql_list.append(row['SQL'])
            g_example_info_list.append(row['资料'])
            g_example_fa_list.append(row['FA'])
            tokens = TOKENIZER(row['问题'])['input_ids']
            g_example_token_list.append(tokens)
    except FileNotFoundError:
        print(f"未找到示例数据文件: {configFinRAG.sql_examples_path}")
    except Exception as e:
        print(f"读取示例数据时出错: {e}")

# 主检索链函数
def sql_retrieve_chain(query):
    if not g_example_question_list:
        load_example_data()
    try:
        result_prompt, sql = generate_sql(query, LLM, g_example_question_list, g_example_sql_list, g_example_token_list)
        # 修正路径中的拼写错误
        db_path = os.path.join('C:', 'Users', '张朋飞', 'Desktop', '数字人', '数字人', 'Fay-fay-agent-edition0830', 'agent', 'tools', 'data', '博金杯比赛数据.db')
        with sqlite3.connect(db_path) as conn:
            cs = conn.cursor()
            success_flag, exc_result = query_db(sql, cs)
        answer = generate_answer(query, exc_result, LLM, g_example_question_list, g_example_info_list, g_example_fa_list,
                                 g_example_token_list)
        return answer
    except sqlite3.Error as e:
        print(f"数据库操作出错: {e}")
    except Exception as e:
        print(f"检索链执行出错: {e}")
    return None